# Define your item pipelines here
#
# Don't forget to add your pipeline to the ITEM_PIPELINES setting
# See: https://docs.scrapy.org/en/latest/topics/item-pipeline.html


# useful for handling different item types with a single interface
from itemadapter import ItemAdapter


from itemadapter import ItemAdapter
from redis import StrictRedis
from scrapy.exceptions import DropItem
import pymongo


class CheckexistPipeline:
    """
    检查数据是否已经存在
    """

    def __init__(self, redis_host, redis_port, redis_database):
        self.redis_host = redis_host
        self.redis_port = redis_port
        self.redis_database = redis_database


    @classmethod
    def from_crawler(cls, crawler):
        redis_host = crawler.settings.get('REDIS_HOST')
        redis_port = crawler.settings.get('REDIS_PORT')
        redis_database = crawler.settings.get('REDIS_DATABASE')
        return cls(redis_host, redis_port, redis_database)

    def open_spider(self, spider):
        try:
            self.redis = StrictRedis(
                host=self.redis_host,
                port=self.redis_port,
                db=self.redis_database,
                decode_responses=True
            )
            self.pipe = self.redis.pipeline()
        except ConnectionError as e:
            spider.logger.error(f"Redis连接失败: {str(e)}")
            raise

    def close_spider(self, spider):
        if hasattr(self, 'pipe'):
            self.pipe.execute()  # 执行剩余命令
        self.redis.close()

    def process_item(self, item, spider):
        hash_val = item.get("hash_val")
        bvid = item.get("bvid")
        if not all([hash_val, bvid]):
            raise DropItem("Invalid item missing required fields")

        redis_key = f"danmaku:{bvid}"
        if self.redis.sismember(redis_key, hash_val):
            raise DropItem(f"Duplicate danmaku: {hash_val}")

        self.pipe.sadd(redis_key, hash_val)
        if len(self.pipe) >= 10:  # 每10条执行一次
            self.pipe.execute()
        return item


class SaveMongoPipeline:
    """
    保存到MongoDB
    """

    def __init__(self, mongo_uri, mongo_db, mongo_collection):
        self.mongo_uri = mongo_uri
        self.mongo_db = mongo_db
        self.mongo_collection = mongo_collection
        self.batch_size = 10  # 批量插入大小
        self.buffer = []

    @classmethod
    def from_crawler(cls, crawler):
        return cls(
            mongo_uri=crawler.settings.get('MONGO_URI', 'mongodb://localhost:27017'),
            mongo_db=crawler.settings.get('MONGO_DATABASE', 'bilibili'),
            mongo_collection=crawler.settings.get('MONGO_COLLECTION', 'danmaku')
        )

    def open_spider(self, spider):
        try:
            self.client = pymongo.MongoClient(self.mongo_uri)
            self.db = self.client[self.mongo_db]
            self.collection = self.db[self.mongo_collection]
            self.collection.create_index([('bvid', pymongo.ASCENDING), ('hash_val', pymongo.ASCENDING)], unique=True)
        except Exception as e:
            spider.logger.error(f"Failed to connect to MongoDB: {e}")

    def close_spider(self, spider):
        if self.buffer:
            self._insert_batch(spider)
        self.client.close()

    def _insert_batch(self, spider):
        try:
            self.collection.insert_many(self.buffer, ordered=False)
        except Exception as e:
            spider.logger.error(f"数据插入失败: {str(e)}")
        finally:
            self.buffer.clear()

    def process_item(self, item, spider):
        self.buffer.append(dict(item))
        if len(self.buffer) >= self.batch_size:
            self._insert_batch(spider)
        return item

